#!/bin/bash -u
# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

function check_exit_code() {
  if (($? != 0 && $? != 124)); then
    exit 1
  fi
}

################################################################################
################################################################################
## Make directory to store fuzzer output
################################################################################
################################################################################
if [ ! -d "out" ]; then
  mkdir -p out
fi

################################################################################
################################################################################
## Set Fuzzer Options
################################################################################
################################################################################
FUZZER_BIN="$SRC/AFL/afl-fuzz"
FUZZER_OPTIONS=""

# Set timeout (if any)
if [[ -n ${TIMEOUT_MS-} ]]; then
  FUZZER_OPTIONS="$FUZZER_OPTIONS -t $TIMEOUT_MS"
fi

# Set memory limit (if any)
if [[ -n ${MEMORY_LIMIT_MB-} ]]; then
  FUZZER_OPTIONS="$FUZZER_OPTIONS -m $MEMORY_LIMIT_MB"
fi

################################################################################
################################################################################
## Fuzz the CORE
################################################################################
################################################################################
if [[ ${INTERACTIVE_MODE} -ne 0 ]]; then
  # Launch fuzzer in interactive mode (good for debugging)
  echo "Launching fuzzer in INTERACTIVE mode..." | tee -a logs/exp.log

  # Set fuzzer mode (Master or Secondary)
  if [[ $MODE == "m" ]]; then
    FUZZER_OPTIONS="$FUZZER_OPTIONS -M"
  else
    FUZZER_OPTIONS="$FUZZER_OPTIONS -S"
  fi

  # Launch fuzzer
  if [[ -z ${DURATION_MINS-} ]]; then
    $FUZZER_BIN \
      -i seeds \
      -o out \
      $FUZZER_OPTIONS \
      afl_${MODE}_interactive \
      $BIN_DIR/$TOPLEVEL
  else
    timeout --foreground --preserve-status ${DURATION_MINS}m $FUZZER_BIN \
      -i seeds \
      -o out \
      $FUZZER_OPTIONS \
      afl_${MODE}_interactive \
      $BIN_DIR/$TOPLEVEL
  fi
  check_exit_code
else

  # Launch fuzzer in non-interactive mode (and potentially many instances)
  echo "Launching fuzzer in NON-INTERACTIVE mode..." | tee -a logs/exp.log
  for ((num = 1; num <= $NUM_INSTANCES; num++)); do

    # Set fuzzer mode (Master or Secondary)
    if [[ $num -eq 1 && $MODE == "m" ]]; then
      FUZZER_OPTIONS="$FUZZER_OPTIONS -M"
    else
      FUZZER_OPTIONS="$FUZZER_OPTIONS -S"
    fi

    # Set log filename
    AFL_LOG="logs/afl_${MODE}_${num}.log"

    # Launch fuzzer
    if [[ -z ${DURATION_MINS-} ]]; then
      $FUZZER_BIN \
        -i seeds \
        -o out \
        $FUZZER_OPTIONS \
        afl_${MODE}_${num} \
        $BIN_DIR/$TOPLEVEL \
        2>&1 | tee -a $AFL_LOG
    else
      timeout --foreground --preserve-status ${DURATION_MINS}m $FUZZER_BIN \
        -i seeds \
        -o out \
        $FUZZER_OPTIONS \
        afl_${MODE}_${num} \
        $BIN_DIR/$TOPLEVEL \
        2>&1 | tee -a $AFL_LOG
    fi
    check_exit_code

    # Set process ID of recently launched fuzzer
    FUZZER_PID=$!
    echo "Launched fuzzer instance $num (PID: $FUZZER_PID)" |
      tee -a logs/exp.log

    # Trap SIGINT to kill all fuzzer processes on ctrl-c
    trap "kill $FUZZER_PID" SIGINT
  done

  # Wait for all fuzzers to complete
  wait
fi

# Change permission on all files to allow access outside of container
find out -type d -exec chmod 777 {} \;
find out -type f -exec chmod 644 {} \;

################################################################################
################################################################################
## Done
################################################################################
################################################################################
echo -e "\e[1;32mFUZZING SUCCESSFUL -- Done!\e[0m" | tee -a logs/exp.log
exit 0
